# 导入包
from matplotlib import pyplot as plt
import pandas as pd
import statsmodels.formula.api as smf
from sqlalchemy import create_engine
import pymysql
import statsmodels.api as sm
import numpy as np
import seaborn as sns

# 数据库配置
db_config = {
    'host': 'localhost',
    'user': 'root',
    'password': 'sjk1234',
    'database': 'msy',
    'port': 3306,
    'charset': 'utf8mb4'   # 添加字符集设置
}

# 创建数据库连接引擎
engine = create_engine(
    f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}?charset={db_config['charset']}"
)

# 使用 SQLAlchemy 的 engine 加载数据
chunk_size = 10000
query = """
    SELECT d.* 
    FROM date_1 d 
    WHERE d.trade_date BETWEEN '2023-01-01' AND '2023-12-31' 
      AND d.ts_code = '000001.SZ'
"""
df_chunks = pd.read_sql_query(query, engine, chunksize=chunk_size)
df1 = pd.concat(df_chunks, ignore_index=True)

# 打印列名以确认实际列名
print("DataFrame 的列名:", df1.columns)

# 新增股票涨跌列
try:
    # 假设实际列名为 'closes' 和 'vol'
    df1['zd_closes'] = round((df1['closes'] - df1['closes'].shift(1)) / df1['closes'].shift(1), 2)
    df1['zs_closes'] = round((df1['closes'] - df1['closes'].shift(1)) / df1['closes'].shift(1), 2)
    df1['zs_vol'] = round((df1['vol'] - df1['vol'].shift(1)) / df1['vol'].shift(1), 2)
except KeyError as e:
    print(f"警告: 缺失列 {e}。请检查数据源中是否包含该列。")
    raise

# 打印新增列后的数据预览
print("新增列后数据预览:")
print(df1.head())

# 处理缺失值数据
df1 = df1.dropna(subset=['zd_closes', 'zs_closes', 'zs_vol'])  # 确保删除缺失值后数据完整
print("处理缺失值后的数据预览:")
print(df1.head())

# 筛选自变量
ex = ['zd_closes', 'id', 'ts_code', 'trade_date', 'the_date', 'opens', 'high', 'low', 'closes', 'pre_closes', 'changes']
number_cols = df1.select_dtypes(include=['number']).columns.tolist()
newlist = [col for col in number_cols if col not in ex]

# 检查筛选后的变量数量
if len(newlist) < 3:
    print(f"警告: 筛选后的自变量数量不足 ({len(newlist)})，无法进行主成分分析。")
else:
    # 选择需要进行主成分分析的变量
    X = df1[newlist]  # 使用筛选出的数值型列

    # 计算特征值和特征向量
    cov_matrix = X.cov()  # 协方差矩阵
    eigenvalues, eigenvectors = np.linalg.eig(cov_matrix)
    print("特征值:", eigenvalues)
    print("特征向量:", eigenvectors)

    # 计算累积贡献率
    total_variance = eigenvalues.sum()
    explained_variance_ratio = eigenvalues / total_variance
    print("累积贡献率:", np.round(explained_variance_ratio.cumsum(), 4))

    # 确保特征向量数量足够
    n_components = min(3, len(eigenvalues))  # 最多取前三个主成分
    top_eigenvectors = eigenvectors[:, :n_components]

    # 计算主成分得分
    principal_components = np.dot(X, top_eigenvectors)
    principal_df = pd.DataFrame(
        principal_components, 
        columns=[f'PC{i+1}' for i in range(n_components)]
    )

    # 合并数据
    data_pca = pd.concat([df1, principal_df], axis=1)

    # 添加常数项
    X_pca = data_pca[[f'PC{i+1}' for i in range(n_components)]].copy()
    X_pca = sm.add_constant(X_pca)

    # 设置因变量
    y = data_pca['zd_closes'].copy()

    # 构建回归模型
    model_pca = sm.OLS(y, X_pca)

    # 拟合模型
    result_pca = model_pca.fit()

    # 输出结果
    print("主成分回归结果:")
    print(result_pca.summary())

    # 提取 PC1、PC2、PC3 作为自变量
    X_pca_selected = data_pca[[f'PC{i+1}' for i in range(n_components)]]
    X_pca_selected = sm.add_constant(X_pca_selected)